from utils.data_helper import LoadSentenceClassificationDataset
from utils.data_helper import tokenizer_en
from config.config import cfg
import os
from loguru import logger

if __name__ == '__main__':
    test_path = os.path.join(cfg.dataset_dir, 'test.csv')
    data_loader = LoadSentenceClassificationDataset(train_file_path=test_path,
                                                    tokenizer=tokenizer_en,
                                                    max_sen_len=None,
                                                    batch_size=cfg.batch_size)

    train_iter, test_iter = data_loader.load_train_val_test_data(test_path, test_path)
    for samble, label in train_iter:
        logger.debug(samble.shape)
        # [seq_len, batch_size]
        logger.debug(samble)
        break
